# !/usr/bin/env python
# -*-coding:utf-8 -*-
"""
# Author     ：hu_cl
# Date       ：2022/2/25 10:15
# File       : listen_proto.py
# Description： 监听 icmp ask asn 和udp
"""
import datetime
import json
import platform
import re
import socket
import struct
import threading
import time
import os
import ctypes
import psycopg2

DATABASES = {
    'default': {
        'HOST': '192.168.2.235',
        'PORT': '5432',
        'NAME': 'hucl_test',
        'USER': 'postgres',
        'PASSWORD': 'sams111',
    }
}

addrs = socket.getaddrinfo(socket.gethostname(), None)
local_ip = ';'.join([item[4][0] for item in addrs if ':' not in item[4][0]])


def system_get_platform_info():
    sysplatform = platform.system().lower().strip()
    python_bit = platform.architecture()[0]
    python_bit_num = re.findall('(\d+)\w*', python_bit)[0]
    return sysplatform, python_bit_num


sys_platform, python_bit_num = system_get_platform_info()


def load_lib(libpath):
    if python_bit_num == '32':
        if sys_platform == "linux":
            libpath = libpath.replace('\\', '/')
            c_lib = ctypes.cdll.LoadLibrary(libpath)
            code, msg = 200, c_lib
        elif sys_platform == "windows":
            c_lib = ctypes.windll.LoadLibrary(libpath)
            code, msg = 200, c_lib
        else:
            code, msg = 0, None
    else:
        code, msg = 0, None

    return code, msg


class LoadLibIoNet:
    path = os.getcwd()
    pathnow = os.path.abspath('io_net.dll')

    def __init__(self):
        self._p_Command = None
        self._p_creatpcap = None
        self._p_openpcap = None
        self.__hand_len = 64
        self.handle = (ctypes.c_char * self.__hand_len)()
        self.loadlib(LoadLibIoNet.pathnow)

    def loadlib(self, lib_path):
        code, c_lib_handle = load_lib(lib_path)

        # int CreatePcap(const char *devName, void *headle, int length)
        self._p_creatpcap = c_lib_handle.pyCreatePcap
        self._p_creatpcap.argtypes = (ctypes.c_char_p, ctypes.c_void_p, ctypes.c_int)
        self._p_creatpcap.restype = ctypes.c_int

        # int OpenPcap(void *obj);
        self._p_openpcap = c_lib_handle.pyOpenPcap
        self._p_openpcap.argtype = ctypes.c_void_p
        self._p_openpcap.restype = ctypes.c_int

        # command msg
        self._p_Command = c_lib_handle.pyCommand
        self._p_Command.argtypes = (
            ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p, ctypes.c_int)
        self._p_Command.restype = ctypes.c_int

    @property
    def p_creatpcap(self):
        return self._p_creatpcap

    @property
    def p_openpcap(self):
        return self._p_openpcap

    @property
    def p_command(self):
        return self._p_Command


class IoNet:
    def __init__(self):
        self.libpcappolicy = LoadLibIoNet()
        self.__hand_len = 64
        self.handle = (ctypes.c_char * self.__hand_len)()
        self.create("ts/sms_policy/")
        self.open()

    def create(self, devName):
        # self.handle = bytes(64) # handle必须不小于64字节空间
        result = self.libpcappolicy.p_creatpcap(devName.encode('utf-8'), ctypes.cast(self.handle, ctypes.c_void_p),
                                                self.__hand_len)
        return result

    def open(self):
        # Result = pcap_lib.OpenPcap(self.handle)
        result = self.libpcappolicy.p_openpcap(ctypes.byref(self.handle))
        return result

    def command(self, type, inBuffer, inLeng, outLeng):
        """
        :param type:  注册类型
        :param inBuffer: 传入字符串
        :param inLeng: 传入长度
        :param outLeng: 输出长度
        :return: 返回码，返回值
        """
        b_inbuffer = inBuffer
        str_by = ctypes.create_string_buffer(inLeng + 4)
        if type != b"ts/CryptoXorD/":
            str_by.value = b_inbuffer
        else:
            str_by = b_inbuffer

        outBuffer = ctypes.create_string_buffer(outLeng)
        # print(f"Command  type:{type}  inBuffer:{inBuffer}")
        result = self.libpcappolicy.p_command(self.handle, type, str_by, inLeng, outBuffer, outLeng)
        return result, outBuffer


meth_c_lib = IoNet()


def de_en_msg(data, leng):
    code, msg = meth_c_lib.command(b"ts/CryptoXorD/", data, leng, leng)  # 消息解密
    return code, msg


class MsgToDb(object):
    def __init__(self):
        self.host = DATABASES['default']['HOST']
        self.username = DATABASES['default']['USER']
        self.password = DATABASES['default']['PASSWORD']
        self.database = DATABASES['default']['NAME']
        self.port = DATABASES['default']['PORT']
        self.__conn__()

    def __conn__(self):
        self.conn = psycopg2.connect(database=self.database, user=self.username, password=self.password, host=self.host,
                                     port=self.port)

    def __execute__(self, sql):
        try:
            self.cursor = self.conn.cursor()
        except:
            self.__conn__()
            self.cursor = self.conn.cursor()
        self.cursor.execute(sql)

    def __close_cursor__(self):
        self.cursor.close()

    def __commit__(self):
        self.conn.commit()

    def __close__(self):
        self.conn.close()


class IcmpScan:

    def __init__(self, timeout=1000):
        self.__id = 55555
        self.timeout = timeout
        self.icmp_socket = self.rawSocket
        self.udp_socket = self.dgramsocket
        self.finished = False

    @property
    def rawSocket(self):
        try:
            Sock = socket.socket(socket.AF_INET, socket.SOCK_RAW, socket.getprotobyname("icmp"))
            Sock.bind((local_ip, 6000))
            # Sock.ioctl(socket.SIO_RCVALL, socket.RCVALL_ON)
            Sock.settimeout(None)
        except Exception as e:
            Sock = self.rawSocket
        return Sock

    @property
    def dgramsocket(self):
        try:
            Sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            Sock.bind(("0.0.0.0", 6000))
        except Exception as e:
            Sock = self.dgramsocket
        return Sock

    def icmp_recv_packet(self):
        while True:
            try:
                recv_packet, addr = self.icmp_socket.recvfrom(1024)
                type, code, checksum, packet_ID, sequence = struct.unpack("bbHHh", recv_packet[20:28])
                print('packet_ID', packet_ID)
                if packet_ID == self.__id:
                    ipHeaderPacket = recv_packet
                    # print('ipHeaderPacket', ipHeaderPacket)
                    source_ip = "%d.%d.%d.%d" % (
                        ipHeaderPacket[12], ipHeaderPacket[13], ipHeaderPacket[14], ipHeaderPacket[15])
                    dst_ip = "%d.%d.%d.%d" % (
                        ipHeaderPacket[16], ipHeaderPacket[17], ipHeaderPacket[18], ipHeaderPacket[19])
                    print('收到的消息记录', source_ip, '--->>>>', dst_ip)
                    # print('ICMP 数据表', recv_packet[28:])
                    rec_data = recv_packet[28:]
                    # print('struct', struct.unpack("!BBHHHBBHII", recv_packet[:20]))
                    # ttl = struct.unpack("!BBHHHBBHII", recv_packet[:20])[5]
                    icmpHeaderPacket = recv_packet[20:28]
                    icmpType = icmpHeaderPacket[0]
                    # 解密信息发送
                    code, data = de_en_msg(rec_data, len(rec_data))
                    if source_ip == '172.16.16.172':
                        print('code.data', code, data.value)
                    if code < 0:
                        print('解密信息错误', data)
                    else:
                        msg = eval(data.value.decode('utf-8'))
                        if dst_ip == local_ip:
                            print('ICMP', icmpType, data.value.decode('utf-8'))
                            self.insert_sqlite({'protocol': "icmp", 'data': msg})
                        else:
                            continue
            except Exception as e:
                print('处理ICMP包，错误', e)
                if self.finished:
                    continue

    def udp_recv_packet(self):
        while True:
            try:
                rec_data, addr = self.udp_socket.recvfrom(1024)
                code, data = de_en_msg(rec_data, len(rec_data))
                if code < 0:
                    print('解密信息错误', data)
                else:
                    msg = eval(data.value.decode('utf-8'))
                    print('UDP', data.value.decode('utf-8'))
                    self.insert_sqlite({'protocol': "udp", 'data': msg})
            except Exception as e:
                print('错误信息', e)
                if self.finished:
                    continue

    def insert_sqlite(self, data):
        protocol = data.get('protocol')
        rec_ip = ';'.join(data['data'].get('ip'))
        dst_ip = ';'.join([item[4][0] for item in addrs if ':' not in item[4][0]])
        helper_uid = data['data'].get('helper_uid')
        timeArray = time.localtime(data['data'].get('time'))
        occ_time = time.strftime("%Y-%m-%d %H:%M:%S", timeArray)
        add_time = datetime.datetime.now()
        msg = json.dumps(data)

        voiaem_log = MsgToDb()
        sql_string = "insert into voiaem (protocol, rec_ip, dst_ip, helper_uid, occ_time, add_time, msg) values('{}', '{}', '{}', '{}', '{}', '{}', '{}');".format(
            protocol, rec_ip, dst_ip, helper_uid, occ_time, add_time, msg)
        voiaem_log.__execute__(sql_string)
        voiaem_log.__close_cursor__()
        voiaem_log.__commit__()

    def start(self):
        t1 = threading.Thread(target=self.icmp_recv_packet, )
        t1.start()
        t2 = threading.Thread(target=self.udp_recv_packet, )
        t2.start()
        self.finished = True
        time.sleep(self.timeout + 1)
        t1.join()
        t2.join()
        self.icmp_socket.close()
        self.udp_socket.close()


if __name__ == '__main__':
    icmp_scan = IcmpScan()
    icmp_scan.start()
    exit(0)
